/*
* Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.
 * MemCache_Hybrid is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
*/
#include "mmc_bm_proxy.h"
#include <algorithm>
#include <numeric>
#include "mmc_logger.h"
#include "mmc_smem_bm_helper.h"
#include "mmc_ptracer.h"

namespace ock {
namespace mmc {
std::map<std::string, MmcRef<MmcBmProxy>> MmcBmProxyFactory::instances_;
std::mutex MmcBmProxyFactory::instanceMutex_;

Result MmcBmProxy::InitBm(const mmc_bm_init_config_t &initConfig, const mmc_bm_create_config_t &createConfig)
{
    std::lock_guard<std::mutex> lock(mutex_);
    if (started_ && handle_ != nullptr) {
        MMC_LOG_INFO("MmcBmProxy " << name_ << " already init");
        return MMC_OK;
    }
    createConfig_ = createConfig;
    MMC_RETURN_ERROR(smem_set_log_level(initConfig.logLevel), "Failed to set smem bm log level");
    if (initConfig.logFunc != nullptr) {
        MMC_RETURN_ERROR(smem_set_extern_logger(initConfig.logFunc), "Failed to set smem bm extern logger");
    }

    smem_bm_config_t config;
    MMC_RETURN_ERROR(smem_bm_config_init(&config), "Failed to init smem bm config");
    config.flags = initConfig.flags;
    config.startConfigStoreServer = false;
    config.hcomTlsConfig = MmcSmemBmHelper::TransSmemTlsConfig(initConfig.hcomTlsConfig);
    config.storeTlsConfig = MmcSmemBmHelper::TransSmemTlsConfig(initConfig.storeTlsConfig);

    // config.hcomUrl is zero-filled, copy only valid chars, and ensure at least one zero at the end.
    std::copy_n(initConfig.hcomUrl.c_str(), std::min(sizeof(config.hcomUrl) - 1, initConfig.hcomUrl.size()),
                config.hcomUrl);

    MMC_RETURN_ERROR(smem_init(0), "Failed to init smem");

    if (smem_bm_init(initConfig.ipPort.c_str(), initConfig.worldSize, initConfig.deviceId, &config) != 0) {
        MMC_LOG_ERROR("Failed to init smem bm");
        smem_uninit();
        return MMC_ERROR;
    }

    bmRankId_ = smem_bm_get_rank_id();

    auto ret = InternalCreateBm(createConfig);
    if (ret != MMC_OK) {
        MMC_LOG_ERROR("Internal create bm failed");
        smem_bm_uninit(0);
        smem_uninit();
        return ret;
    }

    if (smem_bm_join(handle_, 0) != 0) {
        MMC_LOG_ERROR("Failed to join smem bm");
        smem_bm_destroy(handle_);
        smem_bm_uninit(0);
        smem_uninit();
        return MMC_ERROR;
    }

    gvas_[MEDIA_HBM] = smem_bm_ptr_by_mem_type(handle_, SMEM_MEM_TYPE_DEVICE, bmRankId_);
    gvas_[MEDIA_DRAM] = smem_bm_ptr_by_mem_type(handle_, SMEM_MEM_TYPE_HOST, bmRankId_);
    spaces_[MEDIA_HBM] = smem_bm_get_local_mem_size_by_mem_type(handle_, SMEM_MEM_TYPE_DEVICE);
    spaces_[MEDIA_DRAM] = smem_bm_get_local_mem_size_by_mem_type(handle_, SMEM_MEM_TYPE_HOST);
    started_ = true;

    MMC_LOG_INFO("init bm success, rank:" << bmRankId_ << ", worldSize:" << initConfig.worldSize <<
        ", hbm{" << spaces_[MEDIA_HBM] << "}, dram{" << spaces_[MEDIA_DRAM] << "}");
    return MMC_OK;
}

Result MmcBmProxy::InternalCreateBm(const mmc_bm_create_config_t &createConfig)
{
    if (createConfig.localHBMSize > 0 && createConfig.localDRAMSize == 0) {
        mediaType_ = MEDIA_HBM;
    } else if (createConfig.localDRAMSize > 0 && createConfig.localHBMSize == 0) {
        mediaType_ = MEDIA_DRAM;
    } else {
        MMC_LOG_INFO("dram and hbm hybrid pool");
    }

    smem_bm_data_op_type opType = MmcSmemBmHelper::TransSmemBmDataOpType(createConfig.dataOpType);
    if (opType == SMEMB_DATA_OP_BUTT) {
        MMC_LOG_ERROR("MmcBmProxy unknown data op type " << createConfig.dataOpType);
        return MMC_ERROR;
    }
    handle_ = smem_bm_create(createConfig.id, createConfig.memberSize, opType,
        createConfig.localDRAMSize, createConfig.localHBMSize, createConfig.flags);
    if (handle_ == nullptr) {
        MMC_LOG_ERROR("Failed to create smem bm");
        return MMC_ERROR;
    }

    return MMC_OK;
}

void MmcBmProxy::DestroyBm()
{
    std::lock_guard<std::mutex> lock(mutex_);
    if (!started_) {
        MMC_LOG_ERROR("MmcBmProxy (" << name_ << ") is not init");
        return ;
    }

    if (handle_ != nullptr) {
        smem_bm_destroy(handle_);
        handle_ = nullptr;
        std::fill(gvas_, gvas_ + MEDIA_NONE, nullptr);
    }
    smem_bm_uninit(0);
    smem_uninit();
    started_ = false;
    MMC_LOG_INFO("MmcBmProxy (" << name_ << ") is destroyed successfully");
}

MediaType MmcBmProxy::GetMediaType()
{
    // 从上到下选择写入的介质
    for (MediaType type = MEDIA_HBM; type != MEDIA_NONE;) {
        if (GetGva(type) != 0) {
            return type;
        }
        type = MoveDown(type);
    }
    MMC_LOG_ERROR("MmcBmProxy GetMediaType unknown media type");
    return MEDIA_NONE;
}

std::string MmcBmProxy::GetDataOpType() const
{
    return createConfig_.dataOpType;
}

Result MmcBmProxy::Copy(uint64_t srcBmAddr, uint64_t dstBmAddr, uint64_t size, smem_bm_copy_type type)
{
    if (handle_ == nullptr) {
        MMC_LOG_ERROR("Failed to put data to smem bm, handle is null");
        return MMC_ERROR;
    }
    TP_TRACE_BEGIN(TP_SMEM_BM_PUT);
    smem_copy_params params = {(const void*)srcBmAddr, (void*)dstBmAddr, size};
    auto ret = smem_bm_copy(handle_, &params, type, 0);
    TP_TRACE_END(TP_SMEM_BM_PUT, ret);
    return ret;
}

Result MmcBmProxy::Put(const mmc_buffer* buf, uint64_t bmAddr, uint64_t size)
{
    if (handle_ == nullptr) {
        MMC_LOG_ERROR("Failed to put data to smem bm, handle is null");
        return MMC_ERROR;
    }
    if (buf == nullptr) {
        MMC_LOG_ERROR("Failed to put data to smem bm, buf is null");
        return MMC_ERROR;
    }
    smem_bm_copy_type type = buf->type == MEDIA_DRAM ? SMEMB_COPY_H2G : SMEMB_COPY_L2G;
    if (buf->len > size) {
        MMC_LOG_ERROR("Failed to put data to smem bm, buf size : " << buf->len
                                                                   << " is larger than bm block size : " << size);
        return MMC_ERROR;
    }
    TP_TRACE_BEGIN(TP_SMEM_BM_PUT);
    smem_copy_params params = {(void*)(buf->addr + buf->offset), (void*)bmAddr, buf->len};
    auto ret =
            smem_bm_copy(handle_, &params, type, ASYNC_COPY_FLAG);
    TP_TRACE_END(TP_SMEM_BM_PUT, ret);
    return ret;
}

Result MmcBmProxy::Get(const mmc_buffer* buf, uint64_t bmAddr, uint64_t size)
{
    if (handle_ == nullptr) {
        MMC_LOG_ERROR("Failed to get data to smem bm, handle is null");
        return MMC_ERROR;
    }
    if (buf == nullptr) {
        MMC_LOG_ERROR("Failed to get data to smem bm, buf is null");
        return MMC_ERROR;
    }
    smem_bm_copy_type type = buf->type == MEDIA_DRAM ? SMEMB_COPY_G2H : SMEMB_COPY_G2L;
    if (buf->len > size) {
        MMC_LOG_ERROR("Failed to get data to smem bm, buf length: " << buf->len
                                                                    << " not equal data length: " << size);
        return MMC_ERROR;
    }
    TP_TRACE_BEGIN(TP_SMEM_BM_GET);
    smem_copy_params params = {(void*)bmAddr, (void*)(buf->addr + buf->offset), buf->len};
    auto ret =
            smem_bm_copy(handle_, &params, type, ASYNC_COPY_FLAG);
    TP_TRACE_END(TP_SMEM_BM_GET, ret);
    return ret;
}

Result MmcBmProxy::AsyncPut(const MmcBufferArray& bufArr, const MmcMemBlobDesc& blob)
{
    if (handle_ == nullptr) {
        MMC_LOG_ERROR("Failed to get data to smem bm, handle is null");
        return MMC_ERROR;
    }

    if (bufArr.TotalSize() != blob.size_) {
        MMC_LOG_ERROR("Failed to put data to smem bm, total buffer size : " << bufArr.TotalSize()
                << " is not equal to bm block size: " << blob.size_);
        return MMC_ERROR;
    }

    size_t shift = 0;
    for (const auto& buffer : bufArr.Buffers()) {
        auto addr = blob.gva_ + shift;
        MMC_ASSERT_RETURN(addr - shift == blob.gva_, MMC_ERROR);
        MMC_ASSERT_RETURN(blob.size_ >= shift, MMC_ERROR);
        MMC_RETURN_ERROR(Put(&buffer, addr, blob.size_ - shift), "failed put data to smem bm");
        shift += MmcBufSize(buffer);
    }
    return MMC_OK;
}

Result MmcBmProxy::AsyncGet(const MmcBufferArray& bufArr, const MmcMemBlobDesc& blob)
{
    if (handle_ == nullptr) {
        MMC_LOG_ERROR("Failed to get data to smem bm, handle is null");
        return MMC_ERROR;
    }

    if (bufArr.TotalSize() != blob.size_) {
        MMC_LOG_ERROR("Failed to get data from smem bm, total buffer size : " << bufArr.TotalSize()
                << " is not equal to bm block size: " << blob.size_);
        return MMC_ERROR;
    }

    size_t shift = 0;
    for (const auto &buffer : bufArr.Buffers()) {
        auto addr = blob.gva_ + shift;
        MMC_ASSERT_RETURN(addr - shift == blob.gva_, MMC_ERROR);
        MMC_ASSERT_RETURN(blob.size_ >= shift, MMC_ERROR);
        MMC_RETURN_ERROR(Get(&buffer, addr, blob.size_ - shift), "Failed to get data from smem bm");
        shift += MmcBufSize(buffer);
    }
    return MMC_OK;
}

Result MmcBmProxy::BatchPut(const MmcBufferArray& bufArr, const MmcMemBlobDesc& blob)
{
    if (handle_ == nullptr) {
        MMC_LOG_ERROR("Failed to get data to smem bm, handle is null");
        return MMC_ERROR;
    }
    if (bufArr.TotalSize() != blob.size_) {
        MMC_LOG_ERROR("Failed to get data from smem bm, total buffer size : " << bufArr.TotalSize()
            << " is not equal to bm block size: " << blob.size_);
        return MMC_ERROR;
    }
    size_t shift = 0;
    if (bufArr.Buffers().size() > std::numeric_limits<uint32_t>::max()) {
        MMC_LOG_ERROR("buff size is " << bufArr.Buffers().size());
        return MMC_ERROR;
    }
    uint32_t count = bufArr.Buffers().size();
    void *sources[count];
    void *destinations[count];
    uint64_t dataSizes[count];
    smem_bm_copy_type type = bufArr.Buffers()[0].type == MEDIA_DRAM ? SMEMB_COPY_H2G : SMEMB_COPY_L2G;
    for (size_t i = 0; i < count; ++i) {
        auto buf = &bufArr.Buffers()[i];
        sources[i] = reinterpret_cast<void *>(buf->addr + buf->offset);
        destinations[i] = reinterpret_cast<void *>(blob.gva_ + shift);
        dataSizes[i] = buf->len;
        shift += MmcBufSize(*buf);
    }
    smem_batch_copy_params batch_params = {sources, destinations, dataSizes, count};
    return smem_bm_copy_batch(handle_, &batch_params, type, 0);
}

Result MmcBmProxy::BatchGet(const MmcBufferArray& bufArr, const MmcMemBlobDesc& blob)
{
    if (handle_ == nullptr) {
        MMC_LOG_ERROR("Failed to get data to smem bm, handle is null");
        return MMC_ERROR;
    }
    if (bufArr.TotalSize() != blob.size_) {
        MMC_LOG_ERROR("Failed to get data from smem bm, total buffer size : " << bufArr.TotalSize()
            << " is not equal to bm block size: " << blob.size_);
        return MMC_ERROR;
    }
    size_t shift = 0;
    if (bufArr.Buffers().size() > std::numeric_limits<uint32_t>::max()) {
        MMC_LOG_ERROR("buff size is " << bufArr.Buffers().size());
        return MMC_ERROR;
    }
    uint32_t count = bufArr.Buffers().size();
    void *sources[count];
    void *destinations[count];
    uint64_t dataSizes[count];
    smem_bm_copy_type type = bufArr.Buffers()[0].type == MEDIA_DRAM ? SMEMB_COPY_G2H : SMEMB_COPY_G2L;
    for (size_t i = 0; i < count; ++i) {
        auto buf = &bufArr.Buffers()[i];
        destinations[i] = reinterpret_cast<void *>(buf->addr + buf->offset);
        sources[i] = reinterpret_cast<void *>(blob.gva_ + shift);
        dataSizes[i] = buf->len;
        shift += MmcBufSize(*buf);
    }
    smem_batch_copy_params batch_params = {sources, destinations, dataSizes, count};
    return smem_bm_copy_batch(handle_, &batch_params, type, 0);
}

Result MmcBmProxy::BatchDataPut(std::vector<void *> &sources, std::vector<void *> &destinations,
                                const std::vector<uint64_t> &sizes, MediaType localMedia)
{
    if (sources.empty() || sources.size() != destinations.size() || sources.size() != sizes.size()) {
        MMC_LOG_ERROR("Failed data copy, sources:" << sources.size() << ", destinations:" << destinations.size()
                                                   << ", sizes:" << sizes.size());
        return MMC_ERROR;
    }
    if (handle_ == nullptr) {
        MMC_LOG_ERROR("Failed to put data to smem bm, handle is null");
        return MMC_ERROR;
    }
    if (localMedia == MEDIA_NONE) {
        MMC_LOG_ERROR("Failed to put data to smem bm, media:" << localMedia);
        return MMC_ERROR;
    }

    smem_bm_copy_type type = localMedia == MEDIA_DRAM ? SMEMB_COPY_H2G : SMEMB_COPY_L2G;
    smem_batch_copy_params batch_params = {reinterpret_cast<void **>(sources.data()),
                                           reinterpret_cast<void **>(destinations.data()), sizes.data(),
                                           static_cast<uint32_t>(sources.size())};
    uint64_t totalSize = std::accumulate(sizes.begin(), sizes.end(), 0ULL);
    TP_TRACE_BEGIN(TP_MMC_LOCAL_BATCH_PUT);
    auto ret = smem_bm_copy_batch(handle_, &batch_params, type, 0);
    TP_TRACE_END(TP_MMC_LOCAL_BATCH_PUT, ret);
    TP_TRACE_RECORD(TP_MMC_LOCAL_BATCH_PUT_SIZE, totalSize * 1000ULL, 0);
    return ret;
}

Result MmcBmProxy::BatchDataGet(std::vector<void *> &sources, std::vector<void *> &destinations,
                                const std::vector<uint64_t> &sizes, MediaType localMedia)
{
    if (sources.empty() || sources.size() != destinations.size() || sources.size() != sizes.size()) {
        MMC_LOG_ERROR("Failed data copy, sources:" << sources.size() << ", destinations:" << destinations.size()
                                                   << ", sizes:" << sizes.size());
        return MMC_ERROR;
    }
    if (handle_ == nullptr) {
        MMC_LOG_ERROR("Failed to get data to smem bm, handle is null");
        return MMC_ERROR;
    }
    if (localMedia == MEDIA_NONE) {
        MMC_LOG_ERROR("Failed to get data to smem bm, media:" << localMedia);
        return MMC_ERROR;
    }

    smem_bm_copy_type type = localMedia == MEDIA_DRAM ? SMEMB_COPY_G2H : SMEMB_COPY_G2L;
    smem_batch_copy_params batch_params = {reinterpret_cast<void **>(sources.data()),
                                           reinterpret_cast<void **>(destinations.data()), sizes.data(),
                                           static_cast<uint32_t>(sources.size())};
    uint64_t totalSize = std::accumulate(sizes.begin(), sizes.end(), 0ULL);
    TP_TRACE_BEGIN(TP_MMC_LOCAL_BATCH_GET);
    auto ret = smem_bm_copy_batch(handle_, &batch_params, type, 0);
    TP_TRACE_END(TP_MMC_LOCAL_BATCH_GET, ret);
    TP_TRACE_RECORD(TP_MMC_LOCAL_BATCH_GET_SIZE, totalSize * 1000ULL, 0);
    return ret;
}

Result MmcBmProxy::RegisterBuffer(uint64_t addr, uint64_t size)
{
    std::lock_guard<std::mutex> lock(mutex_);
    auto ret = smem_bm_register_user_mem(handle_, addr, size);
    if (ret != MMC_OK) {
        MMC_LOG_ERROR("Failed to register mem,  ret:" << ret);
    }
    return ret;
}

Result MmcBmProxy::CopyWait()
{
    auto ret = smem_bm_wait(handle_);
    if (ret != MMC_OK) {
        MMC_LOG_ERROR("Failed to wait copy task ret:" << ret);
    }
    return ret;
}

}
}